"""
Copyright (c) 2025 Huawei Technologies Co., Ltd.
openUBMC is licensed under Mulan PSL v2.
You can use this software according to the terms and conditions of the Mulan PSL v2.
You may obtain a copy of Mulan PSL v2 at:
        http://license.coscl.org.cn/MulanPSL2
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
See the Mulan PSL v2 for more details.
"""

from typing import Dict, Any

from ..ir import IRContext, IRAnnotation, IRAssign, IRNode, IRNoneNode
from ..extension_base import extension_processor
from ..symbol_meta import Symbol, SymbolKind, TypeInfo, TypeKind
from ..ir.expressions import IRName
from ..ir.declarations import IRFunctionNode, IRClassNode
from ..builtin.builtin_table import BUILTIN_CONFIG
from ..symbol_registry import get_global_registry


@extension_processor(
    "mcpy_type_alias_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_type_alias"}],
    priority=10,
    stage="type_analysis",
    timing="pre",
)
def process_type_alias_annotation(context: IRContext):
    """处理 @mcpy_type_alias 注解

    支持两种使用方式：
    1. 类定义（推荐用于类型）：
       # @mcpy_type_alias: cpp_type=std::int32_t, header=<cstdint>, protocol=Integer
       class int32_t:
           pass

       # 泛型类型：
       # @mcpy_type_alias: cpp_type=mc::future<T>, default_type=mc::future<void>
       class future(Generic[T]):
           pass

    2. 赋值语句（适合别名）：
       # @mcpy_type_alias: cpp_type=char*
       char_ptr = Pointer[char]

    支持的参数：
    - cpp_type: C++ 类型映射（可包含模板参数如<T>）
    - default_type: 无类型参数时的默认类型（用于泛型）
    - header: 头文件（可多个）
    - protocol: 满足的 Protocol（可多个）
    """
    annotation = context.node
    target_node = annotation.parent
    if not isinstance(annotation, IRAnnotation) or not target_node:
        return

    # 分发处理：根据目标节点类型选择处理方式
    if isinstance(target_node, IRClassNode):
        _process_type_alias_class(annotation, target_node)
    elif isinstance(target_node, IRAssign):
        _process_type_alias_assign(annotation, target_node)


def _process_type_alias_class(annotation: IRAnnotation, target_node: IRClassNode):
    """处理类定义上的 @mcpy_type_alias 注解"""
    type_name = target_node.name
    config = _create_type_config(annotation, type_name)
    ir_module = target_node.get_ir_module()
    if not ir_module:
        return

    # 创建类型符号
    # kind统一使用CLASS，具体是否泛型由类型系统根据cpp_type判断
    type_symbol = Symbol(
        name=type_name,
        kind=SymbolKind.CLASS,
        type_info=TypeInfo(
            kind=TypeKind.CLASS,
            name=type_name,
            module=config["module"],
        ),
        module=config["module"],
        is_imported=True,
        import_path=(str(ir_module.module_path)),
        type_config=config,
    )

    # 注册到符号表
    if target_node.template_params:
        registry = ir_module.get_registry()

        # 更新type_info为GENERIC
        type_symbol.type_info = TypeInfo(
            kind=TypeKind.GENERIC,
            name=type_name,
            is_template=True,
            template_params=target_node.template_params,
            min_template_args=len(target_node.template_params),
            max_template_args=len(target_node.template_params),
            module=config["module"],
        )

        # 注册到模块的模板类型表
        registry._template_types[type_name] = type_symbol
        registry._types[type_name] = type_symbol

        global_registry = get_global_registry()
        global_registry._template_types[type_name] = type_symbol
        global_registry._types[type_name] = type_symbol

        # 将符号添加到模块的导入类型
        ir_module.symbol_table._imported_types[type_name] = type_symbol
    else:
        # 普通类型
        ir_module.symbol_table.register_imported_type(type_name, type_symbol)
        ir_module.symbol_table.register_imported_type_config(type_name, config)

    # 添加到类型约束（根据 protocol 自动推断）
    protocols = config.get("protocols", [])
    if protocols:
        for protocol_name in protocols:
            _add_type_to_protocol_constraint(type_name, protocol_name)
    else:
        # 回退到旧的自动推断逻辑
        _add_to_type_constraints(type_name, config)

    # 将IR节点关联到symbol
    type_symbol.ir_node = target_node

    # 判断是否需要保留类IR节点：
    # - 纯类型别名（如int32_t: pass）：无方法，可以删除IR
    # - 带方法的stub类（如future）：有方法，必须保留IR供lookup_symbol查找
    from ..ir.declarations import IRMethodNode, IRFunctionNode

    has_methods = any(
        isinstance(child, (IRMethodNode, IRFunctionNode))
        for child in target_node.get_child_nodes()
    )

    if not has_methods:
        # 纯类型别名，删除IR节点以避免生成无用代码
        ir_module.replace_node(target_node, IRNoneNode())


def _process_type_alias_assign(annotation: IRAnnotation, target_node: IRAssign):
    """处理赋值语句上的 @mcpy_type_alias 注解"""
    # 获取目标名称
    target_expr = target_node.get_target_expr()
    if not target_expr:
        return

    if not isinstance(target_expr, IRName):
        return  # 只处理简单的变量名

    type_name = target_expr.id
    config = _create_type_config(annotation, type_name)
    type_symbol = _create_type_symbol(type_name, config, target_node)

    ir_module = target_node.get_ir_module()
    if not ir_module:
        return

    # 注册到符号表
    ir_module.symbol_table.register_imported_type(type_name, type_symbol)
    ir_module.symbol_table.register_imported_type_config(type_name, config)

    # 添加到类型约束（根据 protocol 自动推断）
    protocols = config.get("protocols", [])
    if protocols:
        for protocol_name in protocols:
            _add_type_to_protocol_constraint(type_name, protocol_name)
    else:
        # 回退到旧的自动推断逻辑
        _add_to_type_constraints(type_name, config)

    # 注解类型已经注册到模块的符号表中，这个赋值节点可以替换掉了
    ir_module.replace_node(target_node, IRNoneNode())


def _create_type_config(annotation: IRAnnotation, type_name: str) -> Dict[str, Any]:
    """创建类型配置"""

    config = {}

    # 从注解参数中获取配置
    config["cpp_type"] = annotation.get_parameter("cpp_type", type_name)
    config["headers"] = annotation.get_parameters("header")

    # 处理 Protocol（支持多个）
    # 支持两种方式：protocol=X,Y 或 satisfies_protocols=X,Y
    protocols = annotation.get_parameters("protocol")
    satisfies_protocols = annotation.get_parameters("satisfies_protocols")

    # 合并两种方式的结果
    all_protocols = []
    if protocols:
        all_protocols.extend(protocols)
    if satisfies_protocols:
        all_protocols.extend(satisfies_protocols)

    if all_protocols:
        config["protocols"] = all_protocols
        # 同时设置satisfies_protocols（兼容builtin配置格式）
        config["satisfies_protocols"] = all_protocols

    # 设置默认值
    config.setdefault("is_type", True)
    module = annotation.get_ir_module()
    config.setdefault("module", (module.module_full_name if module else "stub"))

    # 处理默认值
    cpp_type = config.get("cpp_type", "")
    if "default_value" not in config:
        if "int" in cpp_type:
            config["default_value"] = "0"
        elif "float" in cpp_type or "double" in cpp_type:
            config["default_value"] = "0.0"
        elif "ptr" in cpp_type or "*" in cpp_type:
            config["default_value"] = "nullptr"
        else:
            config["default_value"] = "0"

    return config


def _create_type_symbol(type_name: str, config: Dict[str, Any], target_node) -> Symbol:
    """创建类型符号"""

    ir_module = target_node.get_ir_module()
    module_name = ir_module.module_full_name if ir_module else "stub"

    return Symbol(
        name=type_name,
        kind=SymbolKind.CLASS,
        type_info=TypeInfo(
            kind=TypeKind.CLASS,
            name=type_name,
            module=module_name,
        ),
        module=module_name,
        is_imported=True,
        import_path=(
            str(ir_module.module_path) if ir_module and ir_module.module_path else None
        ),
        type_config=config,  # 设置类型配置
    )


def _add_to_type_constraints(type_name: str, config: Dict[str, Any]) -> None:
    """将新类型动态添加到相关的类型约束中"""

    # 根据类型类别添加到相应的约束中
    type_category = config.get("type_category", "")
    cpp_type = config.get("cpp_type", "")

    # 判断是否为整数类型
    if type_category == "integer" or any(
        x in cpp_type for x in ["int", "size_t", "ptrdiff_t"]
    ):
        _add_to_constraint("Number", type_name)
        _add_to_constraint("Integer", type_name)

        # 根据有符号/无符号添加到相应约束
        if "uint" in cpp_type or cpp_type in ["size_t", "uintptr_t"]:
            _add_to_constraint("UnsignedInteger", type_name)
        else:
            _add_to_constraint("SignedInteger", type_name)

    # 判断是否为浮点类型
    elif type_category == "float" or any(x in cpp_type for x in ["float", "double"]):
        _add_to_constraint("Number", type_name)
        _add_to_constraint("FloatingPoint", type_name)


def _add_to_constraint(constraint_name: str, type_name: str) -> None:
    """添加类型到指定约束

    如果Protocol存在但没有allowed_types字段，会自动创建该字段。
    这允许stub扩展模块动态添加类型到任何已定义的Protocol。
    """

    if constraint_name in BUILTIN_CONFIG:
        # Protocol存在，检查是否有allowed_types
        if "allowed_types" not in BUILTIN_CONFIG[constraint_name]:
            # 没有allowed_types，创建空列表（支持动态扩展）
            BUILTIN_CONFIG[constraint_name]["allowed_types"] = []

        allowed_types = BUILTIN_CONFIG[constraint_name]["allowed_types"]
        if type_name not in allowed_types:
            allowed_types.append(type_name)


def _add_type_to_protocol_constraint(type_name: str, protocol_name: str) -> None:
    """将类型添加到 Protocol 约束

    Args:
        type_name: 类型名称（如 int32_t, result）
        protocol_name: Protocol 名称（如 Integer, Coroutine）
    """
    # 根据 Protocol 名称自动添加到相关约束
    if protocol_name == "Integer":
        _add_to_constraint("Number", type_name)
        _add_to_constraint("Integer", type_name)
        _add_to_constraint("Comparable", type_name)
    elif protocol_name == "Number":
        _add_to_constraint("Number", type_name)
        _add_to_constraint("Comparable", type_name)
    elif protocol_name == "Comparable":
        _add_to_constraint("Comparable", type_name)
    elif protocol_name == "Iterable":
        _add_to_constraint("Iterable", type_name)
    else:
        # 通用处理：直接添加到对应的Protocol约束
        # 如果BUILTIN_CONFIG中存在该Protocol，就添加到它的allowed_types
        _add_to_constraint(protocol_name, type_name)


@extension_processor(
    "mcpy_headers_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_headers"}],
    priority=10,
    stage="type_analysis",
    timing="pre",
)
def process_headers_annotation(context: IRContext):
    """处理 @mcpy_headers 注解"""
    annotation = context.node
    if not isinstance(annotation, IRAnnotation):
        return

    ir_module = context.module
    for header in annotation.keys():
        ir_module.require_include(header)


@extension_processor(
    "mcpy_namespace_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_namespace"}],
    priority=10,
    stage="pre_analysis",
    timing="pre",
)
def process_namespace_annotation(context: IRContext):
    """处理 @mcpy_namespace 注解"""
    annotation = context.node
    if not isinstance(annotation, IRAnnotation) or not annotation.parameters:
        return

    parent = annotation.parent
    if not parent:
        return

    # 设置模块的命名空间
    namespace = annotation.first_key()
    if namespace:
        parent.set_extension_param("namespace", namespace)


@extension_processor(
    "mcpy_emit_template_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_emit_template"}],
    priority=10,
    stage="type_analysis",
    timing="pre",
)
def process_emit_template_annotation(context: IRContext):
    """处理 @mcpy_emit_template 注解

    为函数定义设置自定义的代码生成模板
    """

    annotation = context.node
    target_node = annotation.parent
    if not isinstance(annotation, IRAnnotation) or not target_node:
        return

    if not isinstance(target_node, IRFunctionNode):
        return  # 只处理函数定义

    # 获取模板内容
    emit_template = annotation.first_key()
    if emit_template:
        # 保存模板到函数节点
        target_node.set_extension_param("emit_template", emit_template)


@extension_processor(
    "mcpy_callback_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_callback"}],
    priority=10,
    stage="type_analysis",
    timing="pre",
)
def process_callback_annotation(context: IRContext):
    """处理 @mcpy_callback 注解

    为函数定义设置自定义的回调处理器
    """

    annotation = context.node
    target_node = annotation.parent
    if not isinstance(annotation, IRAnnotation) or not target_node:
        return

    if not isinstance(target_node, IRFunctionNode):
        return  # 只处理函数定义

    # 获取回调路径
    callback_path = annotation.first_key()
    if callback_path:
        # 保存回调路径到函数节点
        target_node.set_extension_param("callback_path", callback_path)


@extension_processor(
    "mcpy_param_override_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_param_override"}],
    priority=10,
    stage="type_analysis",
    timing="pre",
)
def process_param_override_annotation(context: IRContext):
    """处理 @mcpy_param_override 注解

    声明函数参数的实际 C++ 类型，用于参数类型不匹配时的自动转换
    例如：argc: int 但 C++ 期望 int*，注解为 @mcpy_param_override: argc=Pointer[int]
    """

    annotation = context.node
    target_node = annotation.parent
    if not isinstance(annotation, IRAnnotation) or not target_node:
        return

    if not isinstance(target_node, IRFunctionNode):
        return  # 只处理函数定义

    # 解析参数覆盖信息: argc=Pointer[int]
    param_overrides = {}
    for key in annotation.keys():
        override_type = annotation.get_parameter(key)
        if override_type:
            param_overrides[key] = override_type

    # 保存到函数节点
    if param_overrides:
        target_node.set_extension_param("param_overrides", param_overrides)

        # 使父模块的符号缓存失效，以便重新构建包含 param_overrides 的签名
        ir_module = target_node.get_ir_module()
        if ir_module:
            ir_module.invalidate_symbol_cache()


@extension_processor(
    "mcpy_macro_annotation",
    match_rule=[{"node_type": "annotation", "annotation_name": "mcpy_macro"}],
    priority=10,
    stage="type_analysis",
    timing="pre",
)
def process_macro_annotation(context: IRContext):
    """处理 @mcpy_macro 注解

    标记函数为 C++ 宏，需要在代码生成时对参数进行特殊处理（添加括号保护）
    """
    annotation = context.node
    target_node = annotation.parent
    if not isinstance(annotation, IRAnnotation) or not target_node:
        return

    if not isinstance(target_node, IRFunctionNode):
        return  # 只处理函数定义

    # 获取宏标记值（如 True）
    macro_value = annotation.first_key()
    is_macro = macro_value and str(macro_value).lower() in ("true", "1", "yes")
    
    if is_macro:
        # 保存宏标记到函数节点
        target_node.set_extension_param("is_macro", True)
